#ifndef __ASSEMBLE__
#define __ASSEMBLE__

#include "femfunction.h"
#include "matvec.h"
#include "discreteform.h"
#include "quadrature.h"
#include "meshrepartition.h"

typedef struct ASSEMBLE_INFO_
{
    MATRIX *matrix;
    VECTOR *vector;
    DISCRETEFORM *discreteform;

    // 装填用
    INT *row_owner;        // 该行要发送给哪个邻居
    INT *row_owner_index;  // 该行将属于本进程的第几行 or 将是发送给目标进程的第几行
    bool *submat_col_diag; // 单刚矩阵的列属于对角元还是非对角元
    INT *submat_col_index; // 单刚矩阵的列在最后的编号
    INT *col_global_index; // 列的全局编号

    bool if_rhs;                      // 是否有右端项的存在
    INT *send_row_num_displs;         // 要发送给每个邻居几行的偏移量(在下一个变量里)
    INT *send_colnum_each_row_displs; // 要发送给每个邻居的每一行有几列的偏移量(在下一个变量里)
    INT *send_colindex_each_row;      // 要发送给每个邻居的每一行的元的列号
    DOUBLE *send_ent_each_row;        // 要发送给每个邻居的每一行的元 如果有右端项的话会加在发给每个邻居的最后
    INT send_total_ent_num;           // 要发送给每个邻居的元个数 不含右端项

    // 通信用
    INT *send_row_num;           // 要发送给每个邻居的行数
    INT *send_col_num;           // 要发送给每个邻居的列数
    INT *send_col_num_displs;    // 要发送给每个邻居的列数的偏移量
    INT *recv_row_num;           // 要从每个邻居收到的行数
    INT *recv_row_num_displs;    // 要从每个邻居收到的行数的偏移量
    INT *recv_col_num;           // 要从每个邻居收到的列数
    INT *recv_col_num_displs;    // 要从每个邻居收到的列数的偏移量
    INT recv_total_ent_num;      // 要发送给每个邻居的元个数 不含右端项
    DOUBLE *recv_ent_each_row;   // 要从每个邻居收到的每一行的元 如果有右端项的话会加在最后
    INT *recv_row_index;         // 收到的行在本地的编号(自由度编号而不是行的局部编号)
    INT *recv_colnum_each_row;   // 要从每个邻居收到的每一行有几列
    INT *recv_colindex_each_row; // 要从每个邻居收到的每一行的元的列号
} ASSEMBLE_INFO;

/**
 * @brief 完成刚度矩阵从建立到组装的所有步骤
 * @param DiscreteForm 离散变分形式
 * @return 刚度矩阵
 */
void MatrixAssemble(MATRIX **matrix, VECTOR **rhs, DISCRETEFORM *DiscreteForm, DATATYPE type);
MATRIX *MatrixInitial(DISCRETEFORM *DiscreteForm, DATATYPE type);
VECTOR *RHSInitial(DISCRETEFORM *DiscreteForm, DATATYPE type);
void MatrixStructureBuild(MATRIX *matrix, ASSEMBLE_INFO **Assembleinfo, DISCRETEFORM *DiscreteForm);
void SubMatrixAssemble(DISCRETEFORM *DiscreteForm, INT ElemInd, ELEMENT *Elem,
                       BOUNDARYTYPE *ElemBoundaryType, DOUBLE *SubEntries, DOUBLE *SubVec);
void MatrixAddSubMatrix(MATRIX *A, ASSEMBLE_INFO *assembleinfo, DOUBLE *AA, INT n_row, INT *row, INT n_col, INT *col);
void VectorAddSubVector(VECTOR *Vec, ASSEMBLE_INFO *assembleinfo, DOUBLE *SubVec, INT n_row, INT *row);
void MatrixAddRowMatrix(MATRIX *matrix, INT num, INT rowindex, INT *colindex_global, DOUBLE *entries);
INT FindPosition(INT *Vec, INT start, INT end, INT value);
void DirichletBoundaryTreatmen(DISCRETEFORM *DiscreteForm, MATRIX *Matrix, VECTOR *Rhs, ELEMENT *Elem,
                               INT ind, INT numlocalbd, DOUBLE *BoundaryVec, BOUNDARYTYPE *ElemBoundaryType);
void MatrixComplete(MATRIX *matrix);
void MatrixDeleteDirichletBoundary(MATRIX *matrix, DISCRETEFORM *DiscreteForm);
void VectorsAddZeroDirichletBoundary(VECTORS *oldvecs, VECTORS **newvecs, FEMSPACE *femspace);

void ReorderSubMatrixForCurlElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, DOUBLE *SubEntries, DOUBLE *SubVec);
void ReorderBoundaryVecForCurlElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, DOUBLE *BoundaryVec);
void ReorderBoundaryTypeForCurlElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, BOUNDARYTYPE *BoundaryType);
void ReorderLocalProlongColForCurlElem(FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, INT relate, ELEMENT *CoarseElem, DOUBLE *LocalProlong);
void ReorderLocalProlongRowForCurlElem(FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, INT ind, ELEMENT *FinerElem, DOUBLE *LocalProlong);

void ReorderSubMatrixForDivElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, DOUBLE *SubEntries, DOUBLE *SubVec);
void ReorderBoundaryVecForDivElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, DOUBLE *BoundaryVec);
void ReorderBoundaryTypeForDivElem(DISCRETEFORM *DiscreteForm, INT idx_elem, ELEMENT *Elem, BOUNDARYTYPE *BoundaryType);
void ReorderLocalProlongColForDivElem(FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, INT relate, ELEMENT *CoarseElem, DOUBLE *LocalProlong);
void ReorderLocalProlongRowForDivElem(FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, INT ind, ELEMENT *FinerElem, DOUBLE *LocalProlong);

// for multigrid
void ProlongMatrixAssemble(MATRIX **matrix, FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, DATATYPE type);
void RestrictMatrixAssemble(MATRIX **matrix, FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace, DATATYPE type);
void ProduceRelativeElem(ELEMENT *CoarseElem, ELEMENT *FinerElem, ELEMENT *RelativeElem);
void GetLocalProlong(ELEMENT *CoarseElem, BASE *CoarseBase, ELEMENT *FinerElem,
                     BASE *FinerBase, ELEMENT *RelativeElem, DOUBLE *LocalProlong);
void ProlongMatrixSetSubMatrix(MATRIX *A, ASSEMBLE_INFO *assembleinfo, DOUBLE *AA, INT n_row, INT *row, INT n_col, INT *col);
void ProlongDeleteDirichletBoundary(MATRIX *prolong, FEMSPACE *CoarseSpace, FEMSPACE *FinerSpace);

// 下面三个函数需要结合使用, Begin进行第一次组装, 之后每次都可以通过Re重新组装而不重新计算结构和malloc空间和临时空间, End释放空间, 参数都不应该在调用周期内改变
void MatrixBeginAssemble(MATRIX **matrix, VECTOR **rhs, DISCRETEFORM *DiscreteForm, DATATYPE type);
void MatrixReAssemble(MATRIX *matrix, VECTOR *rhs);
void MatrixEndAssemble(MATRIX *matrix, VECTOR *rhs, DISCRETEFORM *DiscreteForm);

// 根据重分布的信息重新排列插值矩阵和有限元空间
void ProlongMatrixReorder(MATRIX **prolong, FEMSPACE **oldspace, MESH *newmesh, BRIDGE *bridge);
void RenewOldGlobalDOF(FEMSPACE *oldspace, FEMSPACE *newspace, BRIDGE *bridge);

#endif